'''
樽海鞘群优化算法
'''
import numpy  as np
import math
import random
import os
import matplotlib.pyplot as plt
import obj_funs
import time

class SSA():
    def __init__(self,n_dim=2, pop_size=10, max_iter=100, lb=[0], ub=[50],obj_func=None,ax=None):
        self.pop = pop_size
        self.n_dim = n_dim
        self.lb, self.ub = np.array(lb) * np.ones(self.n_dim), np.array(ub) * np.ones(self.n_dim)
        self.func = obj_func
        self.max_iter = max_iter
        self.ax = ax  # chart对象，用于3d绘图描点
        self.X = np.random.uniform(low=self.lb, high=self.ub, size=(self.pop, self.n_dim))
        if (self.ax):
            # 初始化标点
            self.ax.chat3dplot(self.X, self.func, 'ax1', 'o', 'black')
        self.Y = [self.func(self.X[i]) for i in range(len(self.X))]
        self.pbest_x = self.X.copy()
        self.pbest_y = [np.inf for i in range(self.pop)]
        self.gbest_x = self.pbest_x.mean(axis=0).reshape(1, -1)
        self.gbest_y = np.inf
        self.gbest_y_hist = []
        self.update_pbest()
        self.update_gbest()

    def update_pbest(self):

        for i in range(len(self.Y)):
            if self.pbest_y[i] > self.Y[i]:
                self.pbest_x[i] = self.X[i]
                self.pbest_y[i] = self.Y[i]

    def update_gbest(self):

        idx_min = self.pbest_y.index(min(self.pbest_y))
        if self.gbest_y > self.pbest_y[idx_min]:
            self.gbest_x = self.X[idx_min, :].copy()
            self.gbest_y = self.pbest_y[idx_min]



    # Function: Updtade Position
    def update_position(self, c1):
        for i in range(0, self.pop):
            if (i <= self.pop / 2): # 领导者比例
                for j in range(0, self.n_dim):
                    c2 = int.from_bytes(os.urandom(8), byteorder="big") / ((1 << 64) - 1)
                    c3 = int.from_bytes(os.urandom(8), byteorder="big") / ((1 << 64) - 1)
                    if (c3 >= 0.5):  # c3 < 0.5
                        try:
                            self.X[i, j] = np.clip((self.gbest_x[0][j] + c1 * ((self.ub[j] - self.lb[j]) * c2 + self.lb[j])), self.lb[j], self.ub[j])
                        except:
                            self.X[i, j] = np.clip((self.gbest_x[j] + c1 * ((self.ub[j] - self.lb[j]) * c2 + self.lb[j])), self.lb[j], self.ub[j])
                    else:
                        try:
                            self.X[i, j] = np.clip((self.gbest_x[0][j] - c1 * ((self.ub[j] - self.lb[j]) * c2 + self.lb[j])), self.lb[j], self.ub[j])
                        except:
                            self.X[i, j] = np.clip((self.gbest_x[j] - c1 * ((self.ub[j] - self.lb[j]) * c2 + self.lb[j])), self.lb[j], self.ub[j])
            else: # 追随者比例
                for j in range(0, self.n_dim):
                    self.X[i, j] = np.clip(((self.X[i - 1, j] + self.X[i, j]) / 2), self.lb[j], self.ub[j])
        self.Y = [self.func(self.X[i]) for i in range(len(self.X))]  # y = f(x) for all particles

    def ssa(self):
        avg_fit = []
        avg_ts = []  # 平均方差
        time_start = time.time()  # 记录迭代寻优开始时间
        for i in range(self.max_iter):
            c1 = 2 * math.exp(-(4 * ((i+1) / self.max_iter)) ** 2)
            self.update_position(c1)
            avg_fit.append(np.mean(self.Y))
            avg_ts.append(np.var(self.Y))
            self.update_pbest()
            self.update_gbest()
            self.gbest_y_hist.append(self.gbest_y)
            # 展示描点图
            if self.ax and i == self.ax.a1:
                # 最终结果化标点
                self.ax.chat3dplot(self.X, self.func, 'ax2', 'o', 'black')

            elif self.ax and i == self.ax.a2:
                self.ax.chat3dplot(self.X, self.func, 'ax3', 'o', 'black')
        self.best_x, self.best_y = self.gbest_x, self.gbest_y
        time_end = time.time()  # 记录迭代结束时间
        print(f'SSA共花费 {time_end - time_start} 秒')
        print('SSA最优适应度', self.best_y)
        print('SSA最优解', self.best_x)
        if self.ax:
            # 最终结果化标点
            self.ax.chat3dplot(self.X, self.func,'ax4', 'o', 'black')
            plt.show()
        return self.gbest_y_hist,avg_fit,avg_ts


# if __name__ == '__main__':
#     n_dim = 30
#     lb = [-5 for i in range(n_dim)]
#     ub = [5 for i in range(n_dim)]
#     demo_func = obj_funs.F1
#     ssa = SSA(n_dim,pop_size=50, max_iter=150, lb=lb, ub=ub, obj_func=demo_func)
#     ssa.run()
#     print('best_x is ', ssa.gbest_x, 'best_y is', ssa.gbest_y)
#     print(f'{demo_func(ssa.gbest_x)}\t{ssa.gbest_x}')
#     plt.plot(ssa.gbest_y_hist)
#     plt.show()



